{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f63df3e4b8d4fdba",
   "metadata": {},
   "source": [
    "# GPT-like multimodal chatbot\n",
    "\n",
    "\n",
    "<img src=\"https://github.com/user-attachments/assets/2ab9b499-7ca2-4ae9-af72-ccc775f30b4e\" width=\"100\" align=\"left\" /> + \n",
    "<img src=\"https://cdn.mos.cms.futurecdn.net/VgGxJABA8DcfAMpPPwdv6a.jpg\" width=\"200\" align=\"center\"/>\n",
    "\n",
    "[https://github.com/apache/burr](https://github.com/apache/burr) by DAGWorks Inc. (YCW23 & StartX).\n",
    "\n",
    "Take🏠:\n",
    "\n",
    " - high level what is Burr\n",
    " - what you can do with Burr (observing state & being able to debug a particular point in time)\n",
    " - watch a walkthrough of this [notebook here](https://youtu.be/hqutVJyd3TI)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d725b4dd413d01e8",
   "metadata": {},
   "source": [
    "## Agentic Problems\n",
    "1. Why did this LLM call fail?\n",
    "2. Oh crap my code broke, why?\n",
    "3. Things went off the rails, but where?\n",
    "4. etc\n",
    "\n",
    "## Monitoring for the win, right?\n",
    "Well but ... monitoring doesn't help you debug & complete your dev loop\n",
    "\n",
    "1. How do I debug that quickly?\n",
    "2. How do I fix the inputs/code, and restart my agent?\n",
    "3. What if my agent was 20+ steps in … do I have to restart from step 0? or can I go to a specific point in time?\n",
    "\n",
    "\n",
    "## Solution: Burr\n",
    "(Complements our other framework [Hamilton](https://github.com/apache/hamilton))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3422aed8-323f-45dd-ade1-b1f0906c35f3",
   "metadata": {},
   "source": [
    "## 1. Agent application is modeled as State + Actions --> Graph\n",
    "Straightforward multi-modal example below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "7953fe759ae0f510",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-08-05T19:18:03.050318Z",
     "start_time": "2024-08-05T19:18:01.630987Z"
    }
   },
   "outputs": [],
   "source": [
    "import copy\n",
    "\n",
    "from IPython.display import Image, display\n",
    "from IPython.core.display import HTML \n",
    "import openai\n",
    "\n",
    "from burr.core import ApplicationBuilder, State, default, graph, when\n",
    "from burr.core.action import action\n",
    "from burr.tracking import LocalTrackingClient\n",
    "\n",
    "MODES = {\n",
    "    \"answer_question\": \"text\",\n",
    "    \"generate_image\": \"image\",\n",
    "    \"generate_code\": \"code\",\n",
    "    \"unknown\": \"text\",\n",
    "}\n",
    "\n",
    "\n",
    "@action(reads=[], writes=[\"chat_history\", \"prompt\"])\n",
    "def process_prompt(state: State, prompt: str) -> State:\n",
    "    result = {\"chat_item\": {\"role\": \"user\", \"content\": prompt, \"type\": \"text\"}}\n",
    "    state = state.append(chat_history=result[\"chat_item\"])\n",
    "    state = state.update(prompt=prompt)\n",
    "    return state\n",
    "\n",
    "\n",
    "@action(reads=[\"prompt\"], writes=[\"mode\"])\n",
    "def choose_mode(state: State) -> State:\n",
    "    prompt = (\n",
    "        f\"You are a chatbot. You've been prompted this: {state['prompt']}. \"\n",
    "        f\"You have the capability of responding in the following modes: {', '.join(MODES)}. \"\n",
    "        \"Please respond with *only* a single word representing the mode that most accurately \"\n",
    "        \"corresponds to the prompt. Fr instance, if the prompt is 'draw a picture of a cat', \"\n",
    "        \"the mode would be 'generate_image'. If the prompt is \"\n",
    "        \"'what is the capital of France', the mode would be 'answer_question'.\"\n",
    "        \"If none of these modes apply, please respond with 'unknown'.\"\n",
    "    )\n",
    "\n",
    "    llm_result = openai.Client().chat.completions.create(\n",
    "        model=\"gpt-4\",\n",
    "        messages=[\n",
    "            {\"role\": \"system\", \"content\": \"You are a helpful assistant\"},\n",
    "            {\"role\": \"user\", \"content\": prompt},\n",
    "        ],\n",
    "    )\n",
    "    content = llm_result.choices[0].message.content\n",
    "    mode = content.lower()\n",
    "    if mode not in MODES:\n",
    "        mode = \"unknown\"\n",
    "    result = {\"mode\": mode}\n",
    "    return state.update(**result)\n",
    "\n",
    "\n",
    "@action(reads=[\"prompt\", \"chat_history\"], writes=[\"response\"])\n",
    "def prompt_for_more(state: State) -> State:\n",
    "    result = {\n",
    "        \"response\": {\n",
    "            \"content\": \"None of the response modes I support apply to your question. \"\n",
    "                       \"Please clarify?\",\n",
    "            \"type\": \"text\",\n",
    "            \"role\": \"assistant\",\n",
    "        }\n",
    "    }\n",
    "    return state.update(**result)\n",
    "\n",
    "\n",
    "@action(reads=[\"prompt\", \"chat_history\", \"mode\"], writes=[\"response\"])\n",
    "def chat_response(\n",
    "        state: State, prepend_prompt: str, model: str = \"gpt-3.5-turbo\"\n",
    ") -> State:\n",
    "    \n",
    "    chat_history = copy.deepcopy(state[\"chat_history\"])\n",
    "    chat_history[-1][\"content\"] = f\"{prepend_prompt}: {chat_history[-1]['content']}\"\n",
    "    chat_history_api_format = [\n",
    "        {\n",
    "            \"role\": chat[\"role\"],\n",
    "            \"content\": chat[\"content\"],\n",
    "        }\n",
    "        for chat in chat_history\n",
    "    ]\n",
    "    client = openai.Client()\n",
    "    result = client.chat.completions.create(\n",
    "        model=model,\n",
    "        messages=chat_history_api_format,\n",
    "    )\n",
    "    text_response = result.choices[0].message.content\n",
    "    result = {\"response\": {\"content\": text_response, \"type\": MODES[state[\"mode\"]], \"role\": \"assistant\"}}\n",
    "    return state.update(**result)\n",
    "\n",
    "\n",
    "@action(reads=[\"prompt\", \"chat_history\", \"mode\"], writes=[\"response\"])\n",
    "def image_response(state: State, model: str = \"dall-e-2\") -> State:\n",
    "    \"\"\"Generates an image response to the prompt. Optional save function to save the image to a URL.\"\"\"\n",
    "    # raise ValueError(\"Demo error\")\n",
    "    client = openai.Client()\n",
    "    result = client.images.generate(\n",
    "        model=model, prompt=state[\"prompt\"], size=\"1024x1024\", quality=\"standard\", n=1\n",
    "    )\n",
    "    image_url = result.data[0].url\n",
    "    result = {\"response\": {\"content\": image_url, \"type\": MODES[state[\"mode\"]], \"role\": \"assistant\"}}\n",
    "    return state.update(**result)\n",
    "\n",
    "\n",
    "@action(reads=[\"response\", \"mode\"], writes=[\"chat_history\"])\n",
    "def response(state: State) -> State:\n",
    "    # you'd do something specific here based on prior state\n",
    "    result = {\"chat_item\": state[\"response\"]}\n",
    "    return state.append(chat_history=result[\"chat_item\"])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "add99d139d3d72a0",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-07-16T18:23:11.529415Z",
     "start_time": "2024-07-16T18:23:11.008552Z"
    }
   },
   "outputs": [
    {
     "data": {
      "image/svg+xml": [
       "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
       "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
       " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
       "<!-- Generated by graphviz version 12.0.0 (20240704.0754)\n",
       " -->\n",
       "<!-- Pages: 1 -->\n",
       "<svg width=\"497pt\" height=\"304pt\"\n",
       " viewBox=\"0.00 0.00 497.41 304.00\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
       "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 300)\">\n",
       "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-300 493.41,-300 493.41,4 -4,4\"/>\n",
       "<!-- prompt -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>prompt</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M408,-231C408,-231 376.25,-231 376.25,-231 370.25,-231 364.25,-225 364.25,-219 364.25,-219 364.25,-207 364.25,-207 364.25,-201 370.25,-195 376.25,-195 376.25,-195 408,-195 408,-195 414,-195 420,-201 420,-207 420,-207 420,-219 420,-219 420,-225 414,-231 408,-231\"/>\n",
       "<text text-anchor=\"middle\" x=\"392.12\" y=\"-207.95\" font-family=\"Times,serif\" font-size=\"14.00\">prompt</text>\n",
       "</g>\n",
       "<!-- decide_mode -->\n",
       "<g id=\"node3\" class=\"node\">\n",
       "<title>decide_mode</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M318.5,-166C318.5,-166 253.75,-166 253.75,-166 247.75,-166 241.75,-160 241.75,-154 241.75,-154 241.75,-142 241.75,-142 241.75,-136 247.75,-130 253.75,-130 253.75,-130 318.5,-130 318.5,-130 324.5,-130 330.5,-136 330.5,-142 330.5,-142 330.5,-154 330.5,-154 330.5,-160 324.5,-166 318.5,-166\"/>\n",
       "<text text-anchor=\"middle\" x=\"286.12\" y=\"-142.95\" font-family=\"Times,serif\" font-size=\"14.00\">decide_mode</text>\n",
       "</g>\n",
       "<!-- prompt&#45;&gt;decide_mode -->\n",
       "<g id=\"edge5\" class=\"edge\">\n",
       "<title>prompt&#45;&gt;decide_mode</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M364,-195.28C352.07,-188.19 337.97,-179.82 325.07,-172.15\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"327.18,-169.33 316.8,-167.23 323.6,-175.35 327.18,-169.33\"/>\n",
       "</g>\n",
       "<!-- input__prompt -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>input__prompt</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" cx=\"392.12\" cy=\"-278\" rx=\"62.1\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"392.12\" y=\"-272.95\" font-family=\"Times,serif\" font-size=\"14.00\">input: prompt</text>\n",
       "</g>\n",
       "<!-- input__prompt&#45;&gt;prompt -->\n",
       "<g id=\"edge1\" class=\"edge\">\n",
       "<title>input__prompt&#45;&gt;prompt</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M392.12,-259.78C392.12,-254.37 392.12,-248.24 392.12,-242.26\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"395.63,-242.71 392.13,-232.71 388.63,-242.71 395.63,-242.71\"/>\n",
       "</g>\n",
       "<!-- generate_image -->\n",
       "<g id=\"node4\" class=\"node\">\n",
       "<title>generate_image</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M90.25,-101C90.25,-101 12,-101 12,-101 6,-101 0,-95 0,-89 0,-89 0,-77 0,-77 0,-71 6,-65 12,-65 12,-65 90.25,-65 90.25,-65 96.25,-65 102.25,-71 102.25,-77 102.25,-77 102.25,-89 102.25,-89 102.25,-95 96.25,-101 90.25,-101\"/>\n",
       "<text text-anchor=\"middle\" x=\"51.12\" y=\"-77.95\" font-family=\"Times,serif\" font-size=\"14.00\">generate_image</text>\n",
       "</g>\n",
       "<!-- decide_mode&#45;&gt;generate_image -->\n",
       "<g id=\"edge6\" class=\"edge\">\n",
       "<title>decide_mode&#45;&gt;generate_image</title>\n",
       "<path fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" d=\"M241.38,-132.38C238.59,-131.55 235.82,-130.75 233.12,-130 182.69,-115.94 166.94,-115.1 113.33,-101.5\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"114.5,-98.19 103.94,-99.08 112.75,-104.96 114.5,-98.19\"/>\n",
       "</g>\n",
       "<!-- generate_code -->\n",
       "<g id=\"node6\" class=\"node\">\n",
       "<title>generate_code</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M202.12,-101C202.12,-101 132.12,-101 132.12,-101 126.12,-101 120.12,-95 120.12,-89 120.12,-89 120.12,-77 120.12,-77 120.12,-71 126.12,-65 132.12,-65 132.12,-65 202.12,-65 202.12,-65 208.12,-65 214.12,-71 214.12,-77 214.12,-77 214.12,-89 214.12,-89 214.12,-95 208.12,-101 202.12,-101\"/>\n",
       "<text text-anchor=\"middle\" x=\"167.12\" y=\"-77.95\" font-family=\"Times,serif\" font-size=\"14.00\">generate_code</text>\n",
       "</g>\n",
       "<!-- decide_mode&#45;&gt;generate_code -->\n",
       "<g id=\"edge7\" class=\"edge\">\n",
       "<title>decide_mode&#45;&gt;generate_code</title>\n",
       "<path fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" d=\"M253.29,-129.62C239.9,-122.53 224.22,-114.23 209.95,-106.67\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"211.66,-103.62 201.19,-102.03 208.39,-109.81 211.66,-103.62\"/>\n",
       "</g>\n",
       "<!-- answer_question -->\n",
       "<g id=\"node7\" class=\"node\">\n",
       "<title>answer_question</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M326.5,-101C326.5,-101 243.75,-101 243.75,-101 237.75,-101 231.75,-95 231.75,-89 231.75,-89 231.75,-77 231.75,-77 231.75,-71 237.75,-65 243.75,-65 243.75,-65 326.5,-65 326.5,-65 332.5,-65 338.5,-71 338.5,-77 338.5,-77 338.5,-89 338.5,-89 338.5,-95 332.5,-101 326.5,-101\"/>\n",
       "<text text-anchor=\"middle\" x=\"285.12\" y=\"-77.95\" font-family=\"Times,serif\" font-size=\"14.00\">answer_question</text>\n",
       "</g>\n",
       "<!-- decide_mode&#45;&gt;answer_question -->\n",
       "<g id=\"edge8\" class=\"edge\">\n",
       "<title>decide_mode&#45;&gt;answer_question</title>\n",
       "<path fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" d=\"M285.85,-129.78C285.77,-124.37 285.67,-118.24 285.57,-112.26\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"289.08,-112.65 285.42,-102.71 282.08,-112.76 289.08,-112.65\"/>\n",
       "</g>\n",
       "<!-- prompt_for_more -->\n",
       "<g id=\"node8\" class=\"node\">\n",
       "<title>prompt_for_more</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M457.88,-101C457.88,-101 368.38,-101 368.38,-101 362.38,-101 356.38,-95 356.38,-89 356.38,-89 356.38,-77 356.38,-77 356.38,-71 362.38,-65 368.38,-65 368.38,-65 457.88,-65 457.88,-65 463.88,-65 469.88,-71 469.88,-77 469.88,-77 469.88,-89 469.88,-89 469.88,-95 463.88,-101 457.88,-101\"/>\n",
       "<text text-anchor=\"middle\" x=\"413.12\" y=\"-77.95\" font-family=\"Times,serif\" font-size=\"14.00\">prompt_for_more</text>\n",
       "</g>\n",
       "<!-- decide_mode&#45;&gt;prompt_for_more -->\n",
       "<g id=\"edge9\" class=\"edge\">\n",
       "<title>decide_mode&#45;&gt;prompt_for_more</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M321.16,-129.62C335.74,-122.39 352.84,-113.91 368.31,-106.23\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"369.44,-109.58 376.84,-102 366.33,-103.31 369.44,-109.58\"/>\n",
       "</g>\n",
       "<!-- response -->\n",
       "<g id=\"node9\" class=\"node\">\n",
       "<title>response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M304.75,-36C304.75,-36 265.5,-36 265.5,-36 259.5,-36 253.5,-30 253.5,-24 253.5,-24 253.5,-12 253.5,-12 253.5,-6 259.5,0 265.5,0 265.5,0 304.75,0 304.75,0 310.75,0 316.75,-6 316.75,-12 316.75,-12 316.75,-24 316.75,-24 316.75,-30 310.75,-36 304.75,-36\"/>\n",
       "<text text-anchor=\"middle\" x=\"285.12\" y=\"-12.95\" font-family=\"Times,serif\" font-size=\"14.00\">response</text>\n",
       "</g>\n",
       "<!-- generate_image&#45;&gt;response -->\n",
       "<g id=\"edge10\" class=\"edge\">\n",
       "<title>generate_image&#45;&gt;response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M102.5,-67.38C105.41,-66.57 108.31,-65.77 111.12,-65 155.72,-52.79 206.9,-39.32 242.15,-30.14\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"242.75,-33.6 251.55,-27.69 240.99,-26.82 242.75,-33.6\"/>\n",
       "</g>\n",
       "<!-- input__model -->\n",
       "<g id=\"node5\" class=\"node\">\n",
       "<title>input__model</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" cx=\"165.12\" cy=\"-148\" rx=\"58.52\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"165.12\" y=\"-142.95\" font-family=\"Times,serif\" font-size=\"14.00\">input: model</text>\n",
       "</g>\n",
       "<!-- input__model&#45;&gt;generate_image -->\n",
       "<g id=\"edge2\" class=\"edge\">\n",
       "<title>input__model&#45;&gt;generate_image</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M137.53,-131.75C124.04,-124.29 107.52,-115.16 92.55,-106.89\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"94.26,-103.84 83.82,-102.07 90.88,-109.97 94.26,-103.84\"/>\n",
       "</g>\n",
       "<!-- input__model&#45;&gt;generate_code -->\n",
       "<g id=\"edge3\" class=\"edge\">\n",
       "<title>input__model&#45;&gt;generate_code</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M165.67,-129.78C165.84,-124.37 166.04,-118.24 166.23,-112.26\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"169.71,-112.81 166.53,-102.71 162.72,-112.59 169.71,-112.81\"/>\n",
       "</g>\n",
       "<!-- input__model&#45;&gt;answer_question -->\n",
       "<g id=\"edge4\" class=\"edge\">\n",
       "<title>input__model&#45;&gt;answer_question</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M193.87,-131.91C208.2,-124.38 225.84,-115.13 241.77,-106.76\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"243.38,-109.87 250.6,-102.12 240.12,-103.67 243.38,-109.87\"/>\n",
       "</g>\n",
       "<!-- generate_code&#45;&gt;response -->\n",
       "<g id=\"edge12\" class=\"edge\">\n",
       "<title>generate_code&#45;&gt;response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M199.68,-64.62C213.08,-57.47 228.76,-49.09 243.01,-41.48\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"244.57,-44.62 251.75,-36.82 241.28,-38.44 244.57,-44.62\"/>\n",
       "</g>\n",
       "<!-- answer_question&#45;&gt;response -->\n",
       "<g id=\"edge11\" class=\"edge\">\n",
       "<title>answer_question&#45;&gt;response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M285.12,-64.78C285.12,-59.37 285.12,-53.24 285.12,-47.26\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"288.63,-47.71 285.13,-37.71 281.63,-47.71 288.63,-47.71\"/>\n",
       "</g>\n",
       "<!-- prompt_for_more&#45;&gt;response -->\n",
       "<g id=\"edge13\" class=\"edge\">\n",
       "<title>prompt_for_more&#45;&gt;response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M377.81,-64.62C362,-56.84 343.24,-47.61 326.76,-39.49\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"328.7,-36.55 318.19,-35.27 325.61,-42.83 328.7,-36.55\"/>\n",
       "</g>\n",
       "<!-- response&#45;&gt;prompt -->\n",
       "<g id=\"edge14\" class=\"edge\">\n",
       "<title>response&#45;&gt;prompt</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M317.03,-21.86C366.8,-27.05 459.27,-39.88 479.12,-65 510.58,-104.78 461.19,-157.01 424.93,-187.3\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"422.78,-184.54 417.22,-193.56 427.19,-189.97 422.78,-184.54\"/>\n",
       "</g>\n",
       "</g>\n",
       "</svg>\n"
      ],
      "text/plain": [
       "<graphviz.graphs.Digraph at 0x12b8d55d0>"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Built the graph.\n",
    "base_graph = (\n",
    "    graph.GraphBuilder()\n",
    "    .with_actions(\n",
    "        # these are the \"nodes\" \n",
    "        prompt=process_prompt,\n",
    "        decide_mode=choose_mode,\n",
    "        generate_image=image_response,\n",
    "        generate_code=chat_response.bind(\n",
    "            prepend_prompt=\"Please respond with *only* code and no other text (at all) to the following:\",\n",
    "        ),\n",
    "        answer_question=chat_response.bind(\n",
    "            prepend_prompt=\"Please answer the following question:\",\n",
    "        ),\n",
    "        prompt_for_more=prompt_for_more,\n",
    "        response=response,\n",
    "    )\n",
    "    .with_transitions(\n",
    "        # these are the edges between nodes, based on state.\n",
    "        (\"prompt\", \"decide_mode\", default),\n",
    "        (\"decide_mode\", \"generate_image\", when(mode=\"generate_image\")),\n",
    "        (\"decide_mode\", \"generate_code\", when(mode=\"generate_code\")),\n",
    "        (\"decide_mode\", \"answer_question\", when(mode=\"answer_question\")),\n",
    "        (\"decide_mode\", \"prompt_for_more\", default),\n",
    "        (\n",
    "            [\"generate_image\", \"answer_question\", \"generate_code\", \"prompt_for_more\"],\n",
    "            \"response\",\n",
    "        ),\n",
    "        (\"response\", \"prompt\", default),\n",
    "    )\n",
    "    .build()\n",
    ")\n",
    "base_graph.visualize()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63b26ad97a43444",
   "metadata": {},
   "source": [
    "## 2. Build application --> built in checkpointing & tracking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "4f0bf205944272ff",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-07-16T18:27:12.136291Z",
     "start_time": "2024-07-16T18:27:11.518066Z"
    }
   },
   "outputs": [
    {
     "data": {
      "image/svg+xml": [
       "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
       "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
       " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
       "<!-- Generated by graphviz version 12.0.0 (20240704.0754)\n",
       " -->\n",
       "<!-- Pages: 1 -->\n",
       "<svg width=\"576pt\" height=\"328pt\"\n",
       " viewBox=\"0.00 0.00 575.76 327.50\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
       "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 323.5)\">\n",
       "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-323.5 571.76,-323.5 571.76,4 -4,4\"/>\n",
       "<!-- prompt -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>prompt</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M505.39,-252.5C505.39,-252.5 473.64,-252.5 473.64,-252.5 467.64,-252.5 461.64,-246.5 461.64,-240.5 461.64,-240.5 461.64,-228.5 461.64,-228.5 461.64,-222.5 467.64,-216.5 473.64,-216.5 473.64,-216.5 505.39,-216.5 505.39,-216.5 511.39,-216.5 517.39,-222.5 517.39,-228.5 517.39,-228.5 517.39,-240.5 517.39,-240.5 517.39,-246.5 511.39,-252.5 505.39,-252.5\"/>\n",
       "<text text-anchor=\"middle\" x=\"489.52\" y=\"-229.45\" font-family=\"Times,serif\" font-size=\"14.00\">prompt</text>\n",
       "</g>\n",
       "<!-- decide_mode -->\n",
       "<g id=\"node3\" class=\"node\">\n",
       "<title>decide_mode</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M436.89,-185.5C436.89,-185.5 372.14,-185.5 372.14,-185.5 366.14,-185.5 360.14,-179.5 360.14,-173.5 360.14,-173.5 360.14,-161.5 360.14,-161.5 360.14,-155.5 366.14,-149.5 372.14,-149.5 372.14,-149.5 436.89,-149.5 436.89,-149.5 442.89,-149.5 448.89,-155.5 448.89,-161.5 448.89,-161.5 448.89,-173.5 448.89,-173.5 448.89,-179.5 442.89,-185.5 436.89,-185.5\"/>\n",
       "<text text-anchor=\"middle\" x=\"404.52\" y=\"-162.45\" font-family=\"Times,serif\" font-size=\"14.00\">decide_mode</text>\n",
       "</g>\n",
       "<!-- prompt&#45;&gt;decide_mode -->\n",
       "<g id=\"edge5\" class=\"edge\">\n",
       "<title>prompt&#45;&gt;decide_mode</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M466.74,-216.08C457.4,-208.94 446.43,-200.55 436.33,-192.82\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"438.63,-190.18 428.56,-186.89 434.38,-195.74 438.63,-190.18\"/>\n",
       "</g>\n",
       "<!-- input__prompt -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>input__prompt</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" cx=\"489.52\" cy=\"-301.5\" rx=\"62.1\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"489.52\" y=\"-296.45\" font-family=\"Times,serif\" font-size=\"14.00\">input: prompt</text>\n",
       "</g>\n",
       "<!-- input__prompt&#45;&gt;prompt -->\n",
       "<g id=\"edge1\" class=\"edge\">\n",
       "<title>input__prompt&#45;&gt;prompt</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M489.52,-283.08C489.52,-277.25 489.52,-270.59 489.52,-264.14\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"493.02,-264.48 489.52,-254.48 486.02,-264.48 493.02,-264.48\"/>\n",
       "</g>\n",
       "<!-- generate_image -->\n",
       "<g id=\"node4\" class=\"node\">\n",
       "<title>generate_image</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M130.64,-103C130.64,-103 52.39,-103 52.39,-103 46.39,-103 40.39,-97 40.39,-91 40.39,-91 40.39,-79 40.39,-79 40.39,-73 46.39,-67 52.39,-67 52.39,-67 130.64,-67 130.64,-67 136.64,-67 142.64,-73 142.64,-79 142.64,-79 142.64,-91 142.64,-91 142.64,-97 136.64,-103 130.64,-103\"/>\n",
       "<text text-anchor=\"middle\" x=\"91.52\" y=\"-79.95\" font-family=\"Times,serif\" font-size=\"14.00\">generate_image</text>\n",
       "</g>\n",
       "<!-- decide_mode&#45;&gt;generate_image -->\n",
       "<g id=\"edge6\" class=\"edge\">\n",
       "<title>decide_mode&#45;&gt;generate_image</title>\n",
       "<path fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" d=\"M359.79,-166.05C290.41,-164.61 160.83,-158.64 122.27,-134.5 114.19,-129.44 107.85,-121.39 103.09,-113.3\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"106.28,-111.85 98.57,-104.58 100.06,-115.07 106.28,-111.85\"/>\n",
       "<text text-anchor=\"middle\" x=\"185.14\" y=\"-121.2\" font-family=\"Times,serif\" font-size=\"14.00\">mode=generate_image</text>\n",
       "</g>\n",
       "<!-- generate_code -->\n",
       "<g id=\"node6\" class=\"node\">\n",
       "<title>generate_code</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M278.52,-103C278.52,-103 208.52,-103 208.52,-103 202.52,-103 196.52,-97 196.52,-91 196.52,-91 196.52,-79 196.52,-79 196.52,-73 202.52,-67 208.52,-67 208.52,-67 278.52,-67 278.52,-67 284.52,-67 290.52,-73 290.52,-79 290.52,-79 290.52,-91 290.52,-91 290.52,-97 284.52,-103 278.52,-103\"/>\n",
       "<text text-anchor=\"middle\" x=\"243.52\" y=\"-79.95\" font-family=\"Times,serif\" font-size=\"14.00\">generate_code</text>\n",
       "</g>\n",
       "<!-- decide_mode&#45;&gt;generate_code -->\n",
       "<g id=\"edge7\" class=\"edge\">\n",
       "<title>decide_mode&#45;&gt;generate_code</title>\n",
       "<path fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" d=\"M359.74,-161.19C323.45,-156.06 276.03,-147.17 261.52,-134.5 255.43,-129.18 251.39,-121.61 248.72,-114.02\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"252.16,-113.31 246.11,-104.61 245.41,-115.18 252.16,-113.31\"/>\n",
       "<text text-anchor=\"middle\" x=\"320.02\" y=\"-121.2\" font-family=\"Times,serif\" font-size=\"14.00\">mode=generate_code</text>\n",
       "</g>\n",
       "<!-- answer_question -->\n",
       "<g id=\"node7\" class=\"node\">\n",
       "<title>answer_question</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M402.89,-103C402.89,-103 320.14,-103 320.14,-103 314.14,-103 308.14,-97 308.14,-91 308.14,-91 308.14,-79 308.14,-79 308.14,-73 314.14,-67 320.14,-67 320.14,-67 402.89,-67 402.89,-67 408.89,-67 414.89,-73 414.89,-79 414.89,-79 414.89,-91 414.89,-91 414.89,-97 408.89,-103 402.89,-103\"/>\n",
       "<text text-anchor=\"middle\" x=\"361.52\" y=\"-79.95\" font-family=\"Times,serif\" font-size=\"14.00\">answer_question</text>\n",
       "</g>\n",
       "<!-- decide_mode&#45;&gt;answer_question -->\n",
       "<g id=\"edge8\" class=\"edge\">\n",
       "<title>decide_mode&#45;&gt;answer_question</title>\n",
       "<path fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" d=\"M397.41,-149.12C393.38,-139.74 388.05,-128.04 382.52,-118 381.59,-116.32 380.61,-114.61 379.59,-112.9\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"382.65,-111.2 374.38,-104.59 376.72,-114.92 382.65,-111.2\"/>\n",
       "<text text-anchor=\"middle\" x=\"454.39\" y=\"-121.2\" font-family=\"Times,serif\" font-size=\"14.00\">mode=answer_question</text>\n",
       "</g>\n",
       "<!-- prompt_for_more -->\n",
       "<g id=\"node8\" class=\"node\">\n",
       "<title>prompt_for_more</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M534.27,-103C534.27,-103 444.77,-103 444.77,-103 438.77,-103 432.77,-97 432.77,-91 432.77,-91 432.77,-79 432.77,-79 432.77,-73 438.77,-67 444.77,-67 444.77,-67 534.27,-67 534.27,-67 540.27,-67 546.27,-73 546.27,-79 546.27,-79 546.27,-91 546.27,-91 546.27,-97 540.27,-103 534.27,-103\"/>\n",
       "<text text-anchor=\"middle\" x=\"489.52\" y=\"-79.95\" font-family=\"Times,serif\" font-size=\"14.00\">prompt_for_more</text>\n",
       "</g>\n",
       "<!-- decide_mode&#45;&gt;prompt_for_more -->\n",
       "<g id=\"edge9\" class=\"edge\">\n",
       "<title>decide_mode&#45;&gt;prompt_for_more</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M449.11,-160.23C478.92,-155.02 514.31,-146.51 523.52,-134.5 529.09,-127.23 526.67,-119.14 521.13,-111.69\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"523.84,-109.47 514.52,-104.43 518.67,-114.18 523.84,-109.47\"/>\n",
       "</g>\n",
       "<!-- response -->\n",
       "<g id=\"node9\" class=\"node\">\n",
       "<title>response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M381.14,-36C381.14,-36 341.89,-36 341.89,-36 335.89,-36 329.89,-30 329.89,-24 329.89,-24 329.89,-12 329.89,-12 329.89,-6 335.89,0 341.89,0 341.89,0 381.14,0 381.14,0 387.14,0 393.14,-6 393.14,-12 393.14,-12 393.14,-24 393.14,-24 393.14,-30 387.14,-36 381.14,-36\"/>\n",
       "<text text-anchor=\"middle\" x=\"361.52\" y=\"-12.95\" font-family=\"Times,serif\" font-size=\"14.00\">response</text>\n",
       "</g>\n",
       "<!-- generate_image&#45;&gt;response -->\n",
       "<g id=\"edge10\" class=\"edge\">\n",
       "<title>generate_image&#45;&gt;response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M142.94,-71.62C193.82,-59.37 270.83,-40.83 318.55,-29.34\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"319.35,-32.75 328.25,-27.01 317.71,-25.95 319.35,-32.75\"/>\n",
       "</g>\n",
       "<!-- input__model -->\n",
       "<g id=\"node5\" class=\"node\">\n",
       "<title>input__model</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" cx=\"58.52\" cy=\"-167.5\" rx=\"58.52\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"58.52\" y=\"-162.45\" font-family=\"Times,serif\" font-size=\"14.00\">input: model</text>\n",
       "</g>\n",
       "<!-- input__model&#45;&gt;generate_image -->\n",
       "<g id=\"edge2\" class=\"edge\">\n",
       "<title>input__model&#45;&gt;generate_image</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M49.26,-149.64C45.25,-139.97 42.5,-127.85 47.52,-118 48.79,-115.5 50.36,-113.14 52.15,-110.94\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"54.25,-113.79 58.83,-104.24 49.29,-108.85 54.25,-113.79\"/>\n",
       "</g>\n",
       "<!-- input__model&#45;&gt;generate_code -->\n",
       "<g id=\"edge3\" class=\"edge\">\n",
       "<title>input__model&#45;&gt;generate_code</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M59.1,-149.33C60.4,-138.81 63.82,-125.9 72.52,-118 72.79,-117.75 135.74,-106.02 185.17,-96.83\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"185.76,-100.28 194.95,-95.01 184.48,-93.4 185.76,-100.28\"/>\n",
       "</g>\n",
       "<!-- input__model&#45;&gt;answer_question -->\n",
       "<g id=\"edge4\" class=\"edge\">\n",
       "<title>input__model&#45;&gt;answer_question</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M111.9,-159.9C143.3,-155.01 183.4,-146.98 217.52,-134.5 231.81,-129.27 233.43,-123.76 247.52,-118 250.2,-116.9 272.65,-110.56 296.93,-103.81\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"297.62,-107.25 306.32,-101.2 295.74,-100.5 297.62,-107.25\"/>\n",
       "</g>\n",
       "<!-- generate_code&#45;&gt;response -->\n",
       "<g id=\"edge12\" class=\"edge\">\n",
       "<title>generate_code&#45;&gt;response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M275.14,-66.58C288.89,-59 305.2,-50.02 319.92,-41.91\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"321.41,-45.09 328.48,-37.2 318.03,-38.96 321.41,-45.09\"/>\n",
       "</g>\n",
       "<!-- answer_question&#45;&gt;response -->\n",
       "<g id=\"edge11\" class=\"edge\">\n",
       "<title>answer_question&#45;&gt;response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M361.52,-66.58C361.52,-60.75 361.52,-54.09 361.52,-47.64\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"365.02,-47.98 361.52,-37.98 358.02,-47.98 365.02,-47.98\"/>\n",
       "</g>\n",
       "<!-- prompt_for_more&#45;&gt;response -->\n",
       "<g id=\"edge13\" class=\"edge\">\n",
       "<title>prompt_for_more&#45;&gt;response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M455.22,-66.58C439.27,-58.48 420.17,-48.78 403.38,-40.26\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"405.13,-37.22 394.62,-35.81 401.96,-43.46 405.13,-37.22\"/>\n",
       "</g>\n",
       "<!-- response&#45;&gt;prompt -->\n",
       "<g id=\"edge14\" class=\"edge\">\n",
       "<title>response&#45;&gt;prompt</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M393.44,-22.2C443.23,-27.91 535.74,-41.73 555.52,-67 589.35,-110.23 545.13,-173.01 514.42,-207.83\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"512.11,-205.17 507.98,-214.93 517.29,-209.87 512.11,-205.17\"/>\n",
       "</g>\n",
       "</g>\n",
       "</svg>\n"
      ],
      "text/plain": [
       "<burr.core.application.Application at 0x12afc6b90>"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tracker = LocalTrackingClient(project=\"agent-demo\")\n",
    "app = (\n",
    "    ApplicationBuilder()\n",
    "    .with_graph(base_graph)\n",
    "    .initialize_from(\n",
    "        tracker, \n",
    "        resume_at_next_action=True, \n",
    "        default_state={\"chat_history\": []},\n",
    "        default_entrypoint=\"prompt\",\n",
    "    )\n",
    "    .with_tracker(tracker)  # tracking + checkpointing; one line 🪄.\n",
    "    .build()\n",
    ")\n",
    "app"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4c48e44fbf36f4d",
   "metadata": {},
   "source": [
    "## 3. Comes with a UI\n",
    "View runs in the UI; Let's run the app first."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "48745f4c04b9d28d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Hi, how can I help? what is the capital of France?\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "🤖: The capital of France is Paris.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Hi, how can I help? write hello world in java\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "🤖: ```\n",
      "public class HelloWorld {\n",
      "    public static void main(String[] args) {\n",
      "        System.out.println(\"Hello, World!\");\n",
      "    }\n",
      "}\n",
      "```\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Hi, how can I help? draw a pen\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "********************************************************************************\n",
      "-------------------------------------------------------------------\n",
      "Oh no an error! Need help with Burr?\n",
      "Join our discord and ask for help! https://discord.gg/4FxBMyzW5n\n",
      "-------------------------------------------------------------------\n",
      "> Action: `generate_image` encountered an error!<\n",
      "> State (at time of action):\n",
      "{'__PRIOR_STEP': 'decide_mode',\n",
      " '__SEQUENCE_ID': 10,\n",
      " 'chat_history': \"[{'role': 'user', 'content': 'what is the capital ...\",\n",
      " 'mode': 'generate_image',\n",
      " 'prompt': 'draw a pen',\n",
      " 'response': \"{'content': '```\\\\npublic class HelloWorld {\\\\n    p...\"}\n",
      "> Inputs (at time of action):\n",
      "{'prompt': 'draw a pen'}\n",
      "********************************************************************************\n",
      "Traceback (most recent call last):\n",
      "  File \"/Users/stefankrawczyk/dagworks/burr/burr/core/application.py\", line 534, in _step\n",
      "    result, new_state = _run_single_step_action(\n",
      "  File \"/Users/stefankrawczyk/dagworks/burr/burr/core/application.py\", line 233, in _run_single_step_action\n",
      "    action.run_and_update(state, **inputs), action.name\n",
      "  File \"/Users/stefankrawczyk/dagworks/burr/burr/core/action.py\", line 533, in run_and_update\n",
      "    return self._fn(state, **self._bound_params, **run_kwargs)\n",
      "  File \"/var/folders/gv/q39lb_1s26x7gbyyypqc3dkm0000gn/T/ipykernel_43564/1354917547.py\", line 94, in image_response\n",
      "    raise ValueError(\"Demo error\")\n",
      "ValueError: Demo error\n"
     ]
    },
    {
     "ename": "ValueError",
     "evalue": "Demo error",
     "output_type": "error",
     "traceback": [
      "\u001B[0;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[0;31mValueError\u001B[0m                                Traceback (most recent call last)",
      "Cell \u001B[0;32mIn[12], line 5\u001B[0m\n\u001B[1;32m      3\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mquit\u001B[39m\u001B[38;5;124m\"\u001B[39m \u001B[38;5;241m==\u001B[39m user_input\u001B[38;5;241m.\u001B[39mlower():\n\u001B[1;32m      4\u001B[0m     \u001B[38;5;28;01mbreak\u001B[39;00m\n\u001B[0;32m----> 5\u001B[0m last_action, action_result, app_state \u001B[38;5;241m=\u001B[39m \u001B[43mapp\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mrun\u001B[49m\u001B[43m(\u001B[49m\n\u001B[1;32m      6\u001B[0m \u001B[43m    \u001B[49m\u001B[43mhalt_after\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43m[\u001B[49m\u001B[38;5;124;43m\"\u001B[39;49m\u001B[38;5;124;43mresponse\u001B[39;49m\u001B[38;5;124;43m\"\u001B[39;49m\u001B[43m]\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\n\u001B[1;32m      7\u001B[0m \u001B[43m    \u001B[49m\u001B[43minputs\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43m{\u001B[49m\u001B[38;5;124;43m\"\u001B[39;49m\u001B[38;5;124;43mprompt\u001B[39;49m\u001B[38;5;124;43m\"\u001B[39;49m\u001B[43m:\u001B[49m\u001B[43m \u001B[49m\u001B[43muser_input\u001B[49m\u001B[43m}\u001B[49m\n\u001B[1;32m      8\u001B[0m \u001B[43m\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m      9\u001B[0m last_message \u001B[38;5;241m=\u001B[39m app_state[\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mchat_history\u001B[39m\u001B[38;5;124m\"\u001B[39m][\u001B[38;5;241m-\u001B[39m\u001B[38;5;241m1\u001B[39m]\n\u001B[1;32m     10\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m last_message[\u001B[38;5;124m'\u001B[39m\u001B[38;5;124mtype\u001B[39m\u001B[38;5;124m'\u001B[39m] \u001B[38;5;241m==\u001B[39m \u001B[38;5;124m'\u001B[39m\u001B[38;5;124mimage\u001B[39m\u001B[38;5;124m'\u001B[39m:\n",
      "File \u001B[0;32m~/dagworks/burr/burr/telemetry.py:273\u001B[0m, in \u001B[0;36mcapture_function_usage.<locals>.wrapped_fn\u001B[0;34m(*args, **kwargs)\u001B[0m\n\u001B[1;32m    270\u001B[0m \u001B[38;5;129m@functools\u001B[39m\u001B[38;5;241m.\u001B[39mwraps(call_fn)\n\u001B[1;32m    271\u001B[0m \u001B[38;5;28;01mdef\u001B[39;00m \u001B[38;5;21mwrapped_fn\u001B[39m(\u001B[38;5;241m*\u001B[39margs, \u001B[38;5;241m*\u001B[39m\u001B[38;5;241m*\u001B[39mkwargs):\n\u001B[1;32m    272\u001B[0m     \u001B[38;5;28;01mtry\u001B[39;00m:\n\u001B[0;32m--> 273\u001B[0m         \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[43mcall_fn\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43margs\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43mkwargs\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m    274\u001B[0m     \u001B[38;5;28;01mfinally\u001B[39;00m:\n\u001B[1;32m    275\u001B[0m         \u001B[38;5;28;01mif\u001B[39;00m is_telemetry_enabled():\n",
      "File \u001B[0;32m~/dagworks/burr/burr/core/application.py:878\u001B[0m, in \u001B[0;36mApplication.run\u001B[0;34m(self, halt_before, halt_after, inputs)\u001B[0m\n\u001B[1;32m    876\u001B[0m \u001B[38;5;28;01mwhile\u001B[39;00m \u001B[38;5;28;01mTrue\u001B[39;00m:\n\u001B[1;32m    877\u001B[0m     \u001B[38;5;28;01mtry\u001B[39;00m:\n\u001B[0;32m--> 878\u001B[0m         \u001B[38;5;28;43mnext\u001B[39;49m\u001B[43m(\u001B[49m\u001B[43mgen\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m    879\u001B[0m     \u001B[38;5;28;01mexcept\u001B[39;00m \u001B[38;5;167;01mStopIteration\u001B[39;00m \u001B[38;5;28;01mas\u001B[39;00m e:\n\u001B[1;32m    880\u001B[0m         \u001B[38;5;28;01mreturn\u001B[39;00m e\u001B[38;5;241m.\u001B[39mvalue\n",
      "File \u001B[0;32m~/dagworks/burr/burr/core/application.py:823\u001B[0m, in \u001B[0;36mApplication.iterate\u001B[0;34m(self, halt_before, halt_after, inputs)\u001B[0m\n\u001B[1;32m    820\u001B[0m prior_action: Optional[Action] \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;01mNone\u001B[39;00m\n\u001B[1;32m    821\u001B[0m \u001B[38;5;28;01mwhile\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mhas_next_action():\n\u001B[1;32m    822\u001B[0m     \u001B[38;5;66;03m# self.step will only return None if there is no next action, so we can rely on tuple unpacking\u001B[39;00m\n\u001B[0;32m--> 823\u001B[0m     prior_action, result, state \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mstep\u001B[49m\u001B[43m(\u001B[49m\u001B[43minputs\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43minputs\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m    824\u001B[0m     \u001B[38;5;28;01myield\u001B[39;00m prior_action, result, state\n\u001B[1;32m    825\u001B[0m     \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_should_halt_iterate(halt_before, halt_after, prior_action):\n",
      "File \u001B[0;32m~/dagworks/burr/burr/core/application.py:495\u001B[0m, in \u001B[0;36mApplication.step\u001B[0;34m(self, inputs)\u001B[0m\n\u001B[1;32m    492\u001B[0m \u001B[38;5;66;03m# we need to increment the sequence before we start computing\u001B[39;00m\n\u001B[1;32m    493\u001B[0m \u001B[38;5;66;03m# that way if we're replaying from state, we don't get stuck\u001B[39;00m\n\u001B[1;32m    494\u001B[0m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_increment_sequence_id()\n\u001B[0;32m--> 495\u001B[0m out \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_step\u001B[49m\u001B[43m(\u001B[49m\u001B[43minputs\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43minputs\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43m_run_hooks\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[38;5;28;43;01mTrue\u001B[39;49;00m\u001B[43m)\u001B[49m\n\u001B[1;32m    496\u001B[0m \u001B[38;5;28;01mreturn\u001B[39;00m out\n",
      "File \u001B[0;32m~/dagworks/burr/burr/core/application.py:548\u001B[0m, in \u001B[0;36mApplication._step\u001B[0;34m(self, inputs, _run_hooks)\u001B[0m\n\u001B[1;32m    546\u001B[0m     exc \u001B[38;5;241m=\u001B[39m e\n\u001B[1;32m    547\u001B[0m     logger\u001B[38;5;241m.\u001B[39mexception(_format_BASE_ERROR_MESSAGE(next_action, \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_state, inputs))\n\u001B[0;32m--> 548\u001B[0m     \u001B[38;5;28;01mraise\u001B[39;00m e\n\u001B[1;32m    549\u001B[0m \u001B[38;5;28;01mfinally\u001B[39;00m:\n\u001B[1;32m    550\u001B[0m     \u001B[38;5;28;01mif\u001B[39;00m _run_hooks:\n",
      "File \u001B[0;32m~/dagworks/burr/burr/core/application.py:534\u001B[0m, in \u001B[0;36mApplication._step\u001B[0;34m(self, inputs, _run_hooks)\u001B[0m\n\u001B[1;32m    532\u001B[0m \u001B[38;5;28;01mtry\u001B[39;00m:\n\u001B[1;32m    533\u001B[0m     \u001B[38;5;28;01mif\u001B[39;00m next_action\u001B[38;5;241m.\u001B[39msingle_step:\n\u001B[0;32m--> 534\u001B[0m         result, new_state \u001B[38;5;241m=\u001B[39m \u001B[43m_run_single_step_action\u001B[49m\u001B[43m(\u001B[49m\n\u001B[1;32m    535\u001B[0m \u001B[43m            \u001B[49m\u001B[43mnext_action\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_state\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43maction_inputs\u001B[49m\n\u001B[1;32m    536\u001B[0m \u001B[43m        \u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m    537\u001B[0m     \u001B[38;5;28;01melse\u001B[39;00m:\n\u001B[1;32m    538\u001B[0m         result \u001B[38;5;241m=\u001B[39m _run_function(\n\u001B[1;32m    539\u001B[0m             next_action, \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_state, action_inputs, name\u001B[38;5;241m=\u001B[39mnext_action\u001B[38;5;241m.\u001B[39mname\n\u001B[1;32m    540\u001B[0m         )\n",
      "File \u001B[0;32m~/dagworks/burr/burr/core/application.py:233\u001B[0m, in \u001B[0;36m_run_single_step_action\u001B[0;34m(action, state, inputs)\u001B[0m\n\u001B[1;32m    230\u001B[0m \u001B[38;5;66;03m# TODO -- guard all reads/writes with a subset of the state\u001B[39;00m\n\u001B[1;32m    231\u001B[0m action\u001B[38;5;241m.\u001B[39mvalidate_inputs(inputs)\n\u001B[1;32m    232\u001B[0m result, new_state \u001B[38;5;241m=\u001B[39m _adjust_single_step_output(\n\u001B[0;32m--> 233\u001B[0m     \u001B[43maction\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mrun_and_update\u001B[49m\u001B[43m(\u001B[49m\u001B[43mstate\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43minputs\u001B[49m\u001B[43m)\u001B[49m, action\u001B[38;5;241m.\u001B[39mname\n\u001B[1;32m    234\u001B[0m )\n\u001B[1;32m    235\u001B[0m _validate_result(result, action\u001B[38;5;241m.\u001B[39mname)\n\u001B[1;32m    236\u001B[0m out \u001B[38;5;241m=\u001B[39m result, _state_update(state, new_state)\n",
      "File \u001B[0;32m~/dagworks/burr/burr/core/action.py:533\u001B[0m, in \u001B[0;36mFunctionBasedAction.run_and_update\u001B[0;34m(self, state, **run_kwargs)\u001B[0m\n\u001B[1;32m    532\u001B[0m \u001B[38;5;28;01mdef\u001B[39;00m \u001B[38;5;21mrun_and_update\u001B[39m(\u001B[38;5;28mself\u001B[39m, state: State, \u001B[38;5;241m*\u001B[39m\u001B[38;5;241m*\u001B[39mrun_kwargs) \u001B[38;5;241m-\u001B[39m\u001B[38;5;241m>\u001B[39m \u001B[38;5;28mtuple\u001B[39m[\u001B[38;5;28mdict\u001B[39m, State]:\n\u001B[0;32m--> 533\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_fn\u001B[49m\u001B[43m(\u001B[49m\u001B[43mstate\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_bound_params\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43mrun_kwargs\u001B[49m\u001B[43m)\u001B[49m\n",
      "Cell \u001B[0;32mIn[8], line 94\u001B[0m, in \u001B[0;36mimage_response\u001B[0;34m(state, model)\u001B[0m\n\u001B[1;32m     91\u001B[0m \u001B[38;5;129m@action\u001B[39m(reads\u001B[38;5;241m=\u001B[39m[\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mprompt\u001B[39m\u001B[38;5;124m\"\u001B[39m, \u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mchat_history\u001B[39m\u001B[38;5;124m\"\u001B[39m, \u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mmode\u001B[39m\u001B[38;5;124m\"\u001B[39m], writes\u001B[38;5;241m=\u001B[39m[\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mresponse\u001B[39m\u001B[38;5;124m\"\u001B[39m])\n\u001B[1;32m     92\u001B[0m \u001B[38;5;28;01mdef\u001B[39;00m \u001B[38;5;21mimage_response\u001B[39m(state: State, model: \u001B[38;5;28mstr\u001B[39m \u001B[38;5;241m=\u001B[39m \u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mdall-e-2\u001B[39m\u001B[38;5;124m\"\u001B[39m) \u001B[38;5;241m-\u001B[39m\u001B[38;5;241m>\u001B[39m State:\n\u001B[1;32m     93\u001B[0m \u001B[38;5;250m    \u001B[39m\u001B[38;5;124;03m\"\"\"Generates an image response to the prompt. Optional save function to save the image to a URL.\"\"\"\u001B[39;00m\n\u001B[0;32m---> 94\u001B[0m     \u001B[38;5;28;01mraise\u001B[39;00m \u001B[38;5;167;01mValueError\u001B[39;00m(\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mDemo error\u001B[39m\u001B[38;5;124m\"\u001B[39m)\n\u001B[1;32m     95\u001B[0m     client \u001B[38;5;241m=\u001B[39m openai\u001B[38;5;241m.\u001B[39mClient()\n\u001B[1;32m     96\u001B[0m     result \u001B[38;5;241m=\u001B[39m client\u001B[38;5;241m.\u001B[39mimages\u001B[38;5;241m.\u001B[39mgenerate(\n\u001B[1;32m     97\u001B[0m         model\u001B[38;5;241m=\u001B[39mmodel, prompt\u001B[38;5;241m=\u001B[39mstate[\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mprompt\u001B[39m\u001B[38;5;124m\"\u001B[39m], size\u001B[38;5;241m=\u001B[39m\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124m1024x1024\u001B[39m\u001B[38;5;124m\"\u001B[39m, quality\u001B[38;5;241m=\u001B[39m\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mstandard\u001B[39m\u001B[38;5;124m\"\u001B[39m, n\u001B[38;5;241m=\u001B[39m\u001B[38;5;241m1\u001B[39m\n\u001B[1;32m     98\u001B[0m     )\n",
      "\u001B[0;31mValueError\u001B[0m: Demo error"
     ]
    }
   ],
   "source": [
    "while True:\n",
    "    user_input = input(\"Hi, how can I help?\")\n",
    "    if \"quit\" == user_input.lower():\n",
    "        break\n",
    "    last_action, action_result, app_state = app.run(\n",
    "        halt_after=[\"response\"], \n",
    "        inputs={\"prompt\": user_input}\n",
    "    )\n",
    "    last_message = app_state[\"chat_history\"][-1]\n",
    "    if last_message['type'] == 'image':\n",
    "        display(Image(url=last_message[\"content\"]))\n",
    "    else:\n",
    "        print(f\"🤖: {last_message['content']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5efe2430b34bf8e5",
   "metadata": {},
   "source": [
    "[Open the UI - http://localhost:7241/](http://localhost:7241)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49bfbd6bf99f7b14",
   "metadata": {},
   "source": [
    "## But something broke / I want to debug\n",
    "\n",
    "Use:\n",
    "\n",
    "* Application ID"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "f5588457208604d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "app_id = \"a6d74912-9ad6-42f0-9a18-bc17c5e77eaf\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "b4abab5aa575628d",
   "metadata": {},
   "outputs": [],
   "source": [
    "resumed_app = (\n",
    "    ApplicationBuilder()\n",
    "    .with_graph(base_graph)\n",
    "    .initialize_from(\n",
    "        tracker,\n",
    "        resume_at_next_action=True, \n",
    "        default_state={\"chat_history\": []},\n",
    "        default_entrypoint=\"prompt\",\n",
    "    )\n",
    "    .with_tracker(tracker)\n",
    "    .with_identifiers(app_id=app_id)\n",
    "    .build()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "4b21ade9-051c-453d-b3fd-f191f970c7b4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'role': 'user', 'content': 'what is the capital of France?', 'type': 'text'},\n",
       " {'content': 'The capital of France is Paris.',\n",
       "  'type': 'text',\n",
       "  'role': 'assistant'},\n",
       " {'role': 'user', 'content': 'write hello world in java', 'type': 'text'},\n",
       " {'content': '```\\npublic class HelloWorld {\\n    public static void main(String[] args) {\\n        System.out.println(\"Hello, World!\");\\n    }\\n}\\n```',\n",
       "  'type': 'code',\n",
       "  'role': 'assistant'},\n",
       " {'role': 'user', 'content': 'draw a pen', 'type': 'text'}]"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "resumed_app.state[\"chat_history\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "2186ef09eaff58d9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Hi, how can I help? \n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<img src=\"https://oaidalleapiprodscus.blob.core.windows.net/private/org-UENfV68Ccazb3rdeYDLRWmoF/user-KU9L39oZGkNLkg6ZfSLMc4AV/img-nONYcLQepD3jBslMG4staoEb.png?st=2024-08-05T20%3A52%3A15Z&se=2024-08-05T22%3A52%3A15Z&sp=r&sv=2023-11-03&sr=b&rscd=inline&rsct=image/png&skoid=6aaadede-4fb3-4698-a8f6-684d7786b067&sktid=a48cca56-e6da-484e-a814-9c849652bcb3&skt=2024-08-05T12%3A52%3A48Z&ske=2024-08-06T12%3A52%3A48Z&sks=b&skv=2023-11-03&sig=qsXqOHTMZhItTdL35Nd4Zic1/Q4LHaDg%2BlFG05rExkw%3D\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Hi, how can I help? what is the capital of England?\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "🤖: The capital of England is London.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Hi, how can I help? quit\n"
     ]
    }
   ],
   "source": [
    "while True:\n",
    "    user_input = input(\"Hi, how can I help?\")\n",
    "    if \"quit\" == user_input.lower():\n",
    "        break\n",
    "    last_action, action_result, app_state = resumed_app.run(\n",
    "        halt_after=[\"response\"], \n",
    "        inputs={\"prompt\": user_input}\n",
    "    )\n",
    "    last_message = app_state[\"chat_history\"][-1]\n",
    "    if last_message['type'] == 'image':\n",
    "        display(Image(url=last_message[\"content\"]))\n",
    "    else:\n",
    "        print(f\"🤖: {last_message['content']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6bec09d1e740a30d",
   "metadata": {},
   "source": [
    "## Actually what if I want to go back to a certain point in time?\n",
    "\n",
    "* Fork: Start with state from any checkpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "c008089be309c098",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-07-16T18:49:55.910255Z",
     "start_time": "2024-07-16T18:49:55.906299Z"
    }
   },
   "outputs": [],
   "source": [
    "app_id = \"a6d74912-9ad6-42f0-9a18-bc17c5e77eaf\"\n",
    "sequence_id = 4\n",
    "# partition_key = \"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "a8c5631cfebe8e48",
   "metadata": {},
   "outputs": [],
   "source": [
    "forked_app = (\n",
    "    ApplicationBuilder()\n",
    "    .with_graph(base_graph) # this could be different...\n",
    "    .initialize_from(\n",
    "        tracker,\n",
    "        resume_at_next_action=True, \n",
    "        default_state={\"chat_history\": []},\n",
    "        default_entrypoint=\"prompt\",\n",
    "        fork_from_app_id=app_id,\n",
    "        fork_from_sequence_id=sequence_id,\n",
    "        # fork_from_partition_key=partition_key\n",
    "    )\n",
    "    .with_tracker(tracker)\n",
    "    .build()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "2a725130e0532f0c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-07-16T18:58:04.223560Z",
     "start_time": "2024-07-16T18:58:04.218274Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'role': 'user', 'content': 'what is the capital of France?', 'type': 'text'},\n",
       " {'content': 'The capital of France is Paris.',\n",
       "  'type': 'text',\n",
       "  'role': 'assistant'},\n",
       " {'role': 'user', 'content': 'write hello world in java', 'type': 'text'}]"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# show prior forked state\n",
    "forked_app.state[\"chat_history\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e27a17968c62aa98",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Hi, how can I help? \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "🤖: ```java\n",
      "public class HelloWorld {\n",
      "    public static void main(String[] args) {\n",
      "        System.out.println(\"Hello, World!\");\n",
      "    }\n",
      "}\n",
      "```\n"
     ]
    }
   ],
   "source": [
    "while True:\n",
    "    user_input = input(\"Hi, how can I help?\")\n",
    "    if \"quit\" == user_input.lower():\n",
    "        break\n",
    "    last_action, action_result, app_state = forked_app.run(\n",
    "        halt_after=[\"response\"], \n",
    "        inputs={\"prompt\": user_input}\n",
    "    )\n",
    "    last_message = app_state[\"chat_history\"][-1]\n",
    "    if last_message['type'] == 'image':\n",
    "        display(Image(url=last_message[\"content\"]))\n",
    "    else:\n",
    "        print(f\"🤖: {last_message['content']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3eb81d1f240d926",
   "metadata": {},
   "source": [
    "## Want to know more?\n",
    "\n",
    "[Link to video walking through this notebook](https://youtu.be/hqutVJyd3TI).\n",
    "\n",
    "[https://github.com/apache/burr](https://github.com/apache/burr)\n",
    "<img src=\"burr_qrcode.png\" width=\"125\"/>\n",
    "\n",
    "[Time Travel blog post & video:](https://blog.dagworks.io/p/travel-back-in-time-with-burr)\n",
    "\n",
    "<a href=\"https://www.youtube.com/embed/98vxhIcE6NI?si=oV1BbSUCKa1UvX5P\"><img src=\"https://img.youtube.com/vi/98vxhIcE6NI/0.jpg\"> \n",
    "</img>\n",
    "<a href=\"https://www.youtube.com/embed/98vxhIcE6NI?si=oV1BbSUCKa1UvX5P\">\n",
    "\n",
    "More blogs @ `blog.dagworks.io` e.g. [async & streaming](https://blog.dagworks.io/p/streaming-chatbot-with-burr-fastapi)\n",
    "\n",
    "More [examples](https://github.com/apache/burr/tree/main/examples/):\n",
    "\n",
    "- e.g. [test case creation](https://burr.apache.org/examples/guardrails/creating_tests/)\n",
    "- e.g. [multi-agent collaboration](https://github.com/apache/burr/tree/main/examples/multi-agent-collaboration)\n",
    "\n",
    "Follow on Twitter & LinkedIn:\n",
    "\n",
    "- https://x.com/burr_framework\n",
    "- https://x.com/dagworks\n",
    "- https://x.com/stefkrawczyk\n",
    "- https://www.linkedin.com/in/skrawczyk/"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
